# Copyright 2015 gRPC authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""The Python implementation of the GRPC helloworld.Greeter server."""
import json
from concurrent import futures
import logging
import grpc
import dlink_pb2
import dlink_pb2_grpc
import base64
import datetime
import torch
from ALILS_AFF.models.model_builder import get_model
from ALILS_AFF.opts import get_arguments


class Dlink(dlink_pb2_grpc.DlinkServicer):
    def Payload(self, request, context):
        print(datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S"))
        # print(request.cmd)
        # print(request.data)
        # 算法部分
        args = get_arguments()
        model, criterion, optimizer, device = get_model(
            layers=args.layers,
            num_class=12,
            in_2d_channels=1,
            dropout=args.dropout,
            lr=args.lr,
            trans="RG",
            mode=args.mode,
            input_len=args.input_len,
            att=args.att,
            fusion=args.fusion,
        )
        model.load_state_dict(
            torch.load(
                "apps\\dgiot_atomgit\\priv\\example\\python\\ALILS_AFF\\checkpoint\\whitedRGCNNc.pth"
            )["state_dict"]
        )
        input_data = request.data
        model.eval()
        with torch.no_grad():
            output = model(input_data)
            pre = torch.softmax(output, 1)

        return dlink_pb2.PayloadResponse(topic="dgiot/ddd", payload="ddddd", ack=pre)


def serve():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    dlink_pb2_grpc.add_DlinkServicer_to_server(Dlink(), server)
    server.add_insecure_port("[::]:30051")
    print("start")
    server.start()
    print("start dd")
    server.wait_for_termination()


if __name__ == "__main__":
    logging.basicConfig()
    serve()
